import os
import time
from srt import parse
from pydub import AudioSegment


def process_input(username, file_num):
    audio_file = f"dataset_raw/{username}/{file_num}.wav"
    subtitle_file = f"dataset_raw/{username}/{file_num}.srt"
    output_dir = f"dataset_wavs/{username}/{file_num}"
    train_list_file = f"dataset_wavs/{username}/{file_num}_train_lists.txt"
    return audio_file, subtitle_file, output_dir, train_list_file

def create_output_dir(output_dir):
    if os.path.exists(output_dir):
        print("\033[93mOutput directory exists. Deleting and recreating...\033[0m")
        time.sleep(3)
        os.rmdir(output_dir)
    os.makedirs(output_dir)

def process_audio(audio_file, subtitle_file, output_dir, train_list_file):
    create_output_dir(output_dir)

    with open(subtitle_file, 'r', encoding='utf-8') as subtitle:
        subtitles = list(parse(subtitle.read()))

    audio = AudioSegment.from_wav(audio_file)
    train_list_entries = []

    for idx, subtitle in enumerate(subtitles):
        start_time = subtitle.start.total_seconds() * 1000
        end_time = subtitle.end.total_seconds() * 1000
        audio_slice = audio[start_time:end_time]

        slice_filename = f"{output_dir}/{file_num}_{str(idx + 1).zfill(4)}.wav"
        audio_slice.export(slice_filename, format="wav")

        train_list_entry = f"{slice_filename}|{subtitle.content}\n"
        train_list_entries.append(train_list_entry)

    with open(train_list_file, 'w', encoding='utf-8') as train_list:
        train_list.writelines(train_list_entries)

    print("\033[92mProcessing complete. Output files created.\033[0m")

if __name__ == "__main__":
    file_num = input("Enter file number: ")

    username= "NickYuan"
    audio_file, subtitle_file, output_dir, train_list_file = process_input(username, file_num)

    try:
        process_audio(audio_file, subtitle_file, output_dir, train_list_file)
    except Exception as e:
        print("\033[91mAn error occurred:", str(e), "\033[0m")
